package com.ideaaedi.commonspring.decorator;

import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;

import java.util.Map;

/**
 * 线程装饰器
 *
 * @author <font size = "20" color = "#3CAA3C"><a href="https://gitee.com/JustryDeng">JustryDeng</a></font> <img src="https://gitee.com/JustryDeng/shared-files/raw/master/JustryDeng/avatar.jpg" />
 * @since 1.0.0
 */
@Slf4j
public class MDCTaskDecorator implements TaskDecorator {

    private final TaskDecorator[] taskDecoratorChain;

    public MDCTaskDecorator(@Nullable TaskDecorator... taskDecoratorChain) {
        this.taskDecoratorChain = taskDecoratorChain;
    }

    @Override
    public @NonNull Runnable decorate(@NonNull Runnable runnable) {
        if (taskDecoratorChain == null || taskDecoratorChain.length == 0) {
            return doDecorate(runnable);
        }
        Runnable wrapedRunnable = runnable;
        for (TaskDecorator taskDecorator : taskDecoratorChain) {
            wrapedRunnable = taskDecorator.decorate(wrapedRunnable);
        }
        return doDecorate(runnable);
    }

    /**
     * MDC装饰
     */
    protected Runnable doDecorate(Runnable runnable) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return () -> {
            try {
                if (contextMap != null) {
                    MDC.setContextMap(contextMap);
                }
            } catch (Exception e) {
                log.warn(" MDC.setContextMap(...) occur exception -> {}", e.getMessage());
            }
            runnable.run();
        };
    }
}
